package com.llmops.core;

import java.util.Collections;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class Context {
    private final ConcurrentHashMap<String, Object> variables = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<String, Node.NodeStatus> nodeStatuses = new ConcurrentHashMap<>();

    private final ConcurrentHashMap<String, Map<String, Object>> nodeInputs = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<String, Map<String, Object>> nodeOutputs = new ConcurrentHashMap<>();

    public void setNodeStatus(String nodeId, Node.NodeStatus status) {
        nodeStatuses.put(nodeId, status);
    }

    public Node.NodeStatus getNodeStatus(String nodeId) {
        return nodeStatuses.getOrDefault(nodeId, Node.NodeStatus.PENDING);
    }

    public void setVariable(String key, Object value) {
        variables.put(key, value);
    }

    public <T> T getVariable(String key) {
        return (T) variables.get(key);
    }

    public boolean containsKey(String key) {
        return variables.containsKey(key);
    }

    public void setNodeInput(String nodeId, Map<String, Object> inputs) {
        nodeInputs.put(nodeId, new ConcurrentHashMap<>(inputs));
    }

    public Map<String, Object> getInputParams(String nodeId) {
        return nodeInputs.getOrDefault(nodeId, Collections.emptyMap());
    }

    public void setNodeOutput(String nodeId, Map<String, Object> outputs) {
        nodeOutputs.put(nodeId, new ConcurrentHashMap<>(outputs));
    }

    public Map<String, Object> getOutputParams(String nodeId) {
        return nodeOutputs.getOrDefault(nodeId, Collections.emptyMap());
    }

    public String fillVariable(Context context,String content) {
        String processedPrompt;
        Pattern pattern = Pattern.compile("\\{(.*?)\\}");
        Matcher matcher = pattern.matcher(content);
        StringBuffer sb = new StringBuffer();
        while (matcher.find()) {
            String varName = matcher.group(1);
            Object value = context.getVariable(varName);
            if (value == null) {
                throw new IllegalArgumentException("变量 '" + varName + "' 未在Context中找到");
            }
            matcher.appendReplacement(sb, Matcher.quoteReplacement(value.toString()));
        }
        matcher.appendTail(sb);
        processedPrompt = sb.toString();
        return processedPrompt;
    }
}